#pragma once
#include <iostream>
#include <string>
#include <vector>
#include <stdlib.h>
#include <fstream>
#include <sstream>
#include <windows.h>
#include <shellapi.h>

void TitleOut(){
    system("title ClTask - A Simple Task Quick Kill Manager");
    system("cls");
    std::cout << "+------------------------------+" << std::endl;
    std::cout << "| ClTask - Task Kill Manager   |" << std::endl;
    std::cout << "+------------------------------+" << std::endl;
    return;
}

class recordInfo{
public:
    bool rescue;
};

recordInfo record;

class rwInfo{
public:
    // 变量甲：存储所有带有'['的行
    std::vector<std::string> linesWithBracket;
    // 变量乙：二维字符串，分离后的内容
    std::vector<std::vector<std::string>> separatedContent;

    int ReadFromFile(){
        std::ifstream inputConf("./user_data/config.conf");
        if (!inputConf) {
            std::cerr << "Error opening config file." << std::endl;
            _sleep(1000);
            return -1;
        }
        std::string line;
        linesWithBracket.clear();
        separatedContent.clear();
        while (std::getline(inputConf, line)) {
            if (line == "EOF") break;
            size_t pos = line.find('[');
            if (pos != std::string::npos) {
                linesWithBracket.push_back(line);
                // 分离
                size_t sep = line.find("]= ");
                if (sep != std::string::npos) {
                    std::string key = line.substr(0, sep);
                    std::string value = line.substr(sep + 3); // 3是"]="的长度
                    if (!key.empty() && key[0] == '[') key = key.substr(1); // 去除第一个'['
                    separatedContent.push_back({key, value});
                }
            }
        }
        inputConf.close();
        return 0;
    }
    int WriteToFile(std::string key, std::string content){
        // 确保 user_data 目录存在
        CreateDirectoryA("./user_data", NULL);
        // 检查文件是否存在
        std::ifstream checkFile("./user_data/config.conf");
        bool exists = checkFile.good();
        checkFile.close();
        // 以追加模式写入（如果文件不存在则创建）
        std::ofstream outputConf("./user_data/config.conf", std::ios::app);
        if (!outputConf) {
            std::cerr << "Error opening config file for writing." << std::endl;
            return -1;
        }
        // 如果文件原本不存在，写入初始内容
        if (!exists) {
            outputConf << "# config.conf generated by ClTask" << std::endl;
        }
        outputConf  << "["<< key << "]= " << content << std::endl;
        outputConf.close();
        return 0;
    }
};

rwInfo info;

class rwTask{
public:
    int much = 0;
    std::string taskCall[1024],taskName[1024];
    bool fk[1024],tk[1024];
    int ReadTaskFromFile(){
        std::ifstream inputConf("./user_data/tasks.conf");
        if (!inputConf) {
            std::cerr << "Error opening tasks file." << std::endl;
            _sleep(1000);
            return -1;
        }
        std::string line;
        int i = 0;
        while (std::getline(inputConf, line)) {
            // 跳过注释行和空行
            std::string trimmed = line;
            trimmed.erase(0, trimmed.find_first_not_of(" \t\r\n"));
            trimmed.erase(trimmed.find_last_not_of(" \t\r\n") + 1);
            if (trimmed.empty() || trimmed[0] == '#') continue;
            // 解析一行
            std::istringstream iss(line);
            std::string name, process, restart;
            int force = 0, tree = 0;
            iss >> name >> process >> force >> tree;
            if (!(iss >> restart)) restart = false; // 修复：防止 string 读到末尾时出错
            taskCall[i] = process;
            taskName[i] = name;
            fk[i] = (force != 0);
            tk[i] = (tree != 0);
            i++;
            much++;
        }
        inputConf.close();
        return 0;
    }
    int WriteTaskToFile(){
        // 确保 user_data 目录存在
        CreateDirectoryA("./user_data", NULL);
        std::string askName,askCall;
        bool askTree = false,askForce = false;
        std::cin.ignore();
        std::cout << "Enter task name: ";
        std::getline(std::cin, askName);
        std::cout << "Enter task call: ";
        std::getline(std::cin, askCall);
        std::cout << "Is force? (1/0): ";
        std::cin >> askForce;
        std::cout << "Is tree? (1/0): ";
        std::cin >> askTree;
        std::cout << "按任意键继续...";
        std::cin.ignore(); // 清除输入缓冲区的换行符
        std::cin.get(); // 等待用户输入
        // 检查文件是否存在
        std::ifstream checkFile("./user_data/tasks.conf");
        bool exists = checkFile.good();
        checkFile.close();
        std::ofstream outputConf("./user_data/tasks.conf", std::ios::app);
        if (!outputConf) {
            std::cerr << "Error opening tasks file for writing." << std::endl;
            _sleep(1000);
            return -1;
        }
        // 如果文件原本不存在，写入初始内容
        if (!exists) {
            outputConf << "# tasks.conf generated by ClTask" << std::endl;
            outputConf << "# Name Task Force Tree " << std::endl;
        }
        outputConf << askName << " " << askCall << " "
                   << (askForce ? 1 : 0) << " " << (askTree ? 1 : 0) << std::endl;
        outputConf.close();
        taskCall[++much] = askCall;
        taskName[much] = askName;
        fk[much] = askForce;
        tk[much] = askTree;
        return 0;
    }
};

rwTask tk;

int ShutTask(std::string processName, bool force, bool tree) {
    if(force&tree){
        system(("taskkill /F /T /IM " + processName).c_str());
        return 0;
    }
    if(force){
        system(("taskkill /F /IM " + processName).c_str());
        return 0;
    }
    if(tree){
        system(("taskkill /T /IM " + processName).c_str());
        return 0;
    }
    if(!force&&!tree){
        system(("taskkill /IM " + processName).c_str());
        return 0;
    }
    return -1;
}

int AskKill(){
    for(int i = 0; i < tk.much; i++) {
        std::cout << "Task " << i + 1 << ": " << tk.taskName[i] << std::endl;
    }
    std::cout << "输入对应序号以终止任务: ";
    int choice;
    std::cin >> choice;
    if (choice < 1 || choice > tk.much) {
        std::cerr << "无效的选择。" << std::endl;
        _sleep(1000);
        return -1;
    }
    return ShutTask(tk.taskCall[choice - 1], tk.fk[choice - 1], tk.tk[choice - 1]);
}

int userWant(bool& exitflag){
    TitleOut();
    std::cout << "1. 查看任务列表并终止任务" << std::endl;
    std::cout << "2. 添加新任务" << std::endl;
    std::cout << "3. 添加新配置" << std::endl;
    std::cout << "0. 退出" << std::endl;
    std::string input;
    int choice;
    std::cin >> input;
    if(input.length() != 1 || input[0] > '9' || input[0] < '0'){
        std::cerr << "无效的选择。" << std::endl;
        _sleep(1000);
        return -1;
    }else if(input[0] == '0'){
        exitflag = true;
        return 1;
    }
    choice = input[0] - '0';
    switch (choice) {
        case 1:
            return AskKill();
        case 2:
            return tk.WriteTaskToFile();
        case 3: {
            std::cout << "请输入配置项和值，以空格分离，若误进入请输入\"exit 1\"" << std::endl;
            std::string in_key, in_value;
            std::cin >> in_key >> in_value;
            if(in_key == "exit" && in_value == "1"){
                return 0;
            }
            info.WriteToFile(in_key, in_value);
            return 0;
        }
        case 0:
            return 0;
        default:
            std::cerr << "无效的选择。" << std::endl;
            _sleep(1000);
            return -1;
    }
}

int PushConfigurations(){
    for (const auto& item : info.separatedContent) {
        if (item.size() == 2) {
            const std::string& key = item[0];
            const std::string& value = item[1];
            // 根据键值对执行相应的操作
            if (key == "rescue") {
                record.rescue = (value == "1");
            }
            // 添加更多的键值对处理逻辑
        }
    }
    return 0;
}

bool IsRunAsAdmin()
{
    BOOL isAdmin = FALSE;
    PSID adminGroup = NULL;
    SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
    if (AllocateAndInitializeSid(&NtAuthority, 2,
        SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS,
        0, 0, 0, 0, 0, 0, &adminGroup))
    {
        CheckTokenMembership(NULL, adminGroup, &isAdmin);
        FreeSid(adminGroup);
    }
    return isAdmin;
}

void RequestUAC()
{
    char szPath[MAX_PATH];
    if (GetModuleFileNameA(NULL, szPath, MAX_PATH))
    {
        SHELLEXECUTEINFOA sei = { sizeof(sei) };
        sei.lpVerb = "runas";
        sei.lpFile = szPath;
        sei.hwnd = NULL;
        sei.nShow = SW_NORMAL;
        if (!ShellExecuteExA(&sei))
        {
            DWORD dwError = GetLastError();
            if (dwError == ERROR_CANCELLED)
                MessageBoxA(NULL, "用户拒绝了UAC权限请求。", "提示", MB_OK);
        }
        exit(0); // 退出当前进程
    }
}

// 在 main() 开头调用
int main() {
    if (!IsRunAsAdmin()) {
        RequestUAC();
        return 0;
    }
    system("chcp 65001 > nul");
    std::setlocale(LC_ALL, "zh_CN.UTF-8");
    
    TitleOut();
    tk.ReadTaskFromFile();
    info.ReadFromFile();
    bool exitflag =false;
    while (!exitflag) {
        userWant(exitflag);
    }
    if (record.rescue) {
        int rescueCount = 0;
        std::ifstream rescueItem("./data/rescue.conf");
        if (rescueItem) {
            std::string line;
            while (std::getline(rescueItem, line)) {
                std::cout << line << std::endl;
                rescueCount++;
            }
            rescueItem.close();
            system("cls");
            TitleOut();
            std::cout << "Rescue mode activated." << std::endl;
            system("cmd.exe");
            for (int i = 0; i < rescueCount; i++) {
                std::cout << "Rescue item " << i + 1 << std::endl;
                system(("start " + tk.taskCall[i]).c_str());
            }
        }
    }
    std::cout << "按回车↩︎键继续";
    std::cin.ignore();
    std::cin.get();
    return 0;
}